# indextts_webui.py  —— IndexTTS 语音克隆系统整合版
import os
import sys
import time
import shutil
import threading
import zipfile
import tempfile
import subprocess
from pathlib import Path

current_dir = Path(__file__).resolve().parent
sys.path.append(str(current_dir))
sys.path.append(str(current_dir / "indextts"))

import gradio as gr
from indextts.infer import IndexTTS
from tools.i18n.i18n import I18nAuto

i18n = I18nAuto(language="zh_CN")
os.makedirs("outputs/tasks", exist_ok=True)
os.makedirs("outputs/batch", exist_ok=True)
os.makedirs("prompts", exist_ok=True)

# ---------------- 模型初始化 ----------------
tts = IndexTTS(model_dir="checkpoints", cfg_path="checkpoints/config.yaml")

# ---------------- 推理封装 ----------------
def infer(voice_path: str, text: str, output_path: str | None = None):
    if output_path is None:
        output_path = f"outputs/tasks/spk_{int(time.time()*1000)}.wav"
    with mutex:
        tts.infer(voice_path, text, output_path)
    return output_path

# ---------------- 单条生成 ----------------
def gen_single(prompt, text):
    if not prompt:
        raise gr.Error("请先上传参考音频")
    if not text.strip():
        raise gr.Error("请输入目标文本")
    
    out = infer(prompt, text)
    return gr.update(value=out, visible=True)

def update_prompt_audio():
    return gr.update(interactive=True)

# ---------------- 音频处理 ----------------
def get_audio_duration(audio_path):
    """使用 FFmpeg 获取音频时长（秒）"""
    try:
        # 使用 FFmpeg 获取音频时长
        cmd = ["ffprobe", "-v", "error", "-show_entries", "format=duration", 
               "-of", "default=noprint_wrappers=1:nokey=1", audio_path]
        result = subprocess.run(cmd, capture_output=True, text=True)
        duration = float(result.stdout.strip())
        return duration
    except Exception as e:
        print(f"获取音频时长时出错: {e}")
        return 0  # 出错时返回0

def trim_audio(audio_path, max_seconds=20):
    """截取音频前N秒，如果音频长度小于N秒则保持原样"""
    try:
        # 获取音频时长
        duration = get_audio_duration(audio_path)
        
        # 如果音频长度小于最大秒数，则返回原始路径
        if duration <= max_seconds:
            return audio_path
            
        # 创建临时文件
        temp_dir = Path("cache/temp_audio")
        temp_dir.mkdir(parents=True, exist_ok=True)
        
        # 使用原始文件名但添加前缀
        orig_name = Path(audio_path).name
        temp_path = temp_dir / f"trim_{orig_name}"
        
        # 使用 FFmpeg 截取音频
        cmd = ["ffmpeg", "-y", "-i", audio_path, "-t", str(max_seconds), 
               "-acodec", "copy", str(temp_path)]
        subprocess.run(cmd, capture_output=True)
        
        return str(temp_path)
    except Exception as e:
        print(f"处理音频时出错: {e}")
        return audio_path  # 出错时返回原始路径

# ---------------- 批量生成（多参考音频，一文本）----------------
def scan_prompts():
    """返回 prompts/ 下所有音频文件路径"""
    ext = (".wav", ".flac", ".mp3")
    return [str(p) for p in Path("prompts").rglob("*") if p.suffix.lower() in ext]

def gen_batch(text: str, ref_paths: list, progress=gr.Progress()):
    if not text.strip():
        raise gr.Error("请输入目标文本")
    if not ref_paths:
        raise gr.Error("请至少提供一份参考音频")

    stamp = int(time.time()*1000)
    out_dir = Path(f"outputs/tasks/batch_{stamp}")
    out_dir.mkdir(parents=True, exist_ok=True)
    zip_path = out_dir.with_suffix(".zip")

    # 创建临时目录存放处理后的音频
    temp_dir = Path("cache/temp_audio")
    temp_dir.mkdir(parents=True, exist_ok=True)
    
    progress(0, desc="开始批量生成…")
    for idx, ref in enumerate(ref_paths):
        # 截取音频前20秒
        trimmed_ref = trim_audio(ref, max_seconds=20)
        
        out_wav = out_dir / f"voice{idx+1}.wav"
        infer(trimmed_ref, text, str(out_wav))
        progress((idx+1)/len(ref_paths), desc=f"已完成 {idx+1}/{len(ref_paths)}")

    # 打包
    with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as z:
        for w in out_dir.glob("*.wav"):
            z.write(w, w.name)
    return str(zip_path)

# ---------------- 批量顺序合成（一参考音频，多文本）----------------
def batch_tts(prompt_audio_path, txt_folder_path, progress=gr.Progress()):
    if not prompt_audio_path:
        return "❌ 请先上传参考音频！"
    if not txt_folder_path:
        return "❌ 请输入 txt 文件夹路径！"
    folder = Path(txt_folder_path.strip())
    if not folder.is_dir():
        return f"❌ 路径不存在：{folder}"
    txt_files = sorted(folder.rglob("*.txt"))
    if not txt_files:
        return f"❌ 文件夹 {folder} 内未找到任何 txt 文件！"

    out_dir = Path("outputs/batch")
    out_dir.mkdir(exist_ok=True)
    progress(0, desc="开始批量合成...")
    for idx, txt_path in enumerate(txt_files):
        text = txt_path.read_text(encoding="utf-8").strip()
        if not text:
            continue
        wav_path = out_dir / txt_path.relative_to(folder).with_suffix(".wav")
        wav_path.parent.mkdir(parents=True, exist_ok=True)
        infer(prompt_audio_path, text, str(wav_path))
        progress((idx + 1) / len(txt_files), desc=f"已合成 {idx + 1}/{len(txt_files)}")
    return f"✅ 全部合成完成！结果保存在 {out_dir.resolve()}"

# ---------------- Gradio 界面 ----------------
mutex = threading.Lock()

with gr.Blocks(title="IndexTTS 语音克隆") as demo:
    gr.HTML("""
    <h2><center>IndexTTS: 工业级零样本语音克隆</center></h2>
    <p align="center"><a href='https://www.xiamibaba.com'>Xiamibaba 虾米巴巴</a></p>
    """)

    with gr.Tab("单条生成"):
        with gr.Row():
            prompt_audio = gr.Audio(label="参考音频", sources=["upload", "microphone"], type="filepath")
            input_text_single = gr.Textbox(label="目标文本", lines=3)
            gen_btn = gr.Button("生成", variant="primary")
            out_audio = gr.Audio(label="结果", visible=False)
        
        prompt_audio.upload(update_prompt_audio, inputs=[], outputs=[gen_btn])
        gen_btn.click(gen_single, [prompt_audio, input_text_single], out_audio)

    with gr.Tab("批量生成（多参考音频，一文本）"):
        gr.Markdown("### 步骤：上传/选择多个参考音频 → 输入文本 → 一键打包下载")
        with gr.Row():
            with gr.Column():
                upload_refs = gr.File(label="上传多个参考音频", file_count="multiple", file_types=["audio"])
                scan_btn = gr.Button("扫描 prompts 目录", size="sm")
                ref_listbox = gr.Dropdown(label="已扫描列表（多选）", multiselect=True, interactive=True)
            with gr.Column():
                input_text_batch = gr.Textbox(label="目标文本", lines=4, placeholder="请输入要合成的文本")
                batch_btn = gr.Button("批量生成", variant="primary")
                out_zip = gr.File(label="打包下载", visible=False)

        # 扫描按钮
        scan_btn.click(lambda: gr.update(choices=scan_prompts()), outputs=ref_listbox)

        # 合并上传+扫描，并调用生成
        def gen_batch_wrapper(text, uploaded_files, scanned_files):
            ref_paths = [f.name for f in uploaded_files] if uploaded_files else []
            ref_paths.extend(scanned_files or [])
            return gen_batch(text, ref_paths)

        batch_btn.click(
            fn=gen_batch_wrapper,
            inputs=[input_text_batch, upload_refs, ref_listbox],
            outputs=out_zip
        ).then(lambda: gr.update(visible=True), outputs=out_zip)

    with gr.Tab("批量顺序合成（一参考音频，多文本）"):
        gr.Markdown('### 步骤：1. 先在"单条生成"标签页上传参考音频 2. 在下方输入txt文件夹路径 3. 点击"开始批量合成"')
        with gr.Row():
            folder_path = gr.Textbox(label="txt 文件夹路径", placeholder=r"D:\ai\my_txts")
            batch_seq_btn = gr.Button("开始批量合成", variant="primary")
        progress_tb = gr.Textbox(label="进度/结果", interactive=False, lines=2)
        
        batch_seq_btn.click(batch_tts,
                        inputs=[prompt_audio, folder_path],
                        outputs=progress_tb)

if __name__ == "__main__":
    demo.queue().launch(server_name="127.0.0.1", inbrowser=True, max_threads=5)